from scipy.io import loadmat
from glob import glob
import os
from PIL import Image, ImageDraw, ImageFont
import shutil


def copy_image(src_path, dst_dir):
    """
    将图像文件复制到目标目录（保留文件名）
    
    参数:
        src_path (str): 原图像文件路径
        dst_dir (str): 目标目录路径
    """
    # 确保目标目录存在
    os.makedirs(dst_dir, exist_ok=True)
    
    # 执行复制
    shutil.copy(src_path, dst_dir)
    print(f"✅ 已复制文件到: {os.path.join(dst_dir, os.path.basename(src_path))}")


def recreate_folder(folder_path):
    """
    强制重建文件夹 (如果存在则删除后新建)
    
    参数:
        folder_path (str): 目标文件夹路径
    """
    # 检查是否存在
    if os.path.exists(folder_path):
        # 确保是目录而不是文件
        if os.path.isdir(folder_path):
            # 递归删除整个目录
            shutil.rmtree(folder_path)
            print(f"✅ 已删除旧目录: {folder_path}")
        else:
            raise NotADirectoryError(f"路径 {folder_path} 是文件而非目录")

    # 创建新目录 (自动创建父目录)
    os.makedirs(folder_path, exist_ok=True)
    print(f"✅ 已创建新目录: {folder_path}")

def annotate_image(image_path, points, output_path):
    # 打开图像
    img = Image.open(image_path)
    draw = ImageDraw.Draw(img)
    
    # 设置字体（如果系统没有Arial，可替换其他字体路径）
    try:
        font = ImageFont.truetype("arial.ttf", 15)
    except IOError:
        font = ImageFont.load_default()
    
    # 标注每个点
    for i, (x, y) in enumerate(points, 1):
        radius = 5
        # 绘制红色圆圈
        draw.ellipse([(x-radius, y-radius), (x+radius, y+radius)], 
                     outline="red", width=2)
        # 添加序号
        text_position = (x + radius + 2, y - radius - 2)
        draw.text(text_position, str(i), fill="red", font=font)
    
    # 保存结果
    img.save(output_path)
    print(f"标注后的图像已保存至：{output_path}")



sha_root = 'ShanghaiTech/part_A'
phases = ['test_data', 'train_data']

save_root = 'sha'
recreate_folder(save_root)

for phase in phases:
    phase_ = "test" if 'test' in phase else "train"
    save_path = f"{save_root}/{phase_}"
    recreate_folder(save_path)
    
    image_path = f"{sha_root}/{phase}/images"
    img_list = glob(os.path.join(image_path, '*.jpg'))
    
    file_points = {}
    for img_file in img_list:
        img_name = os.path.basename(img_file).split(".")[0]
        
        gt_file = img_file.replace('images', 'ground-truth').replace('IMG_', 'GT_IMG_').replace('.jpg', '.mat')
        gt_points = loadmat(gt_file)['image_info'][0][0][0][0][0]
        
        file_points[img_file] = []
        for point in gt_points:
            x, y = point
            file_points[img_file].append([x, y])
    
    with open(f'./{save_root}/shanghai_tech_part_a_{phase_}.list', 'w') as f:
        for img_file, points in file_points.items():
            copy_image(img_file, save_path)
            
            img_name = os.path.basename(img_file).split(".")[0]
            label_file = f'{save_path}/{img_name}.txt'
            f.write(f"{phase_}/{img_name}.jpg {phase_}/{img_name}.txt\n")
            with open(label_file, 'w') as pf:
                for point in points:
                    pf.write(f"{point[0]} {point[1]}\n")
